package com.aiwiown.snackmq.storage.timer;

import com.aiwiown.snackmq.common.message.Message;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.NavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 负责定时消息的存储和检索。
 * 它使用一个时间轮（slotTable）和一系列日志段文件（TimerLogSegment）来高效管理定时消息。
 */
@Slf4j
@Getter
public class TimerMessageStore implements Closeable {

    private final String storePath;
    private final int wheelSize;
    private final long segmentSize;

    private final NavigableMap<Long, TimerLogSegment> segments = new ConcurrentSkipListMap<>();
    private final Map<Integer, Integer> slotTable;
    private volatile TimerLogSegment activeSegment;
    private final AtomicLong writePosition = new AtomicLong(0);
    private final ReentrantLock lock = new ReentrantLock();

    /**
     * 用于在内存中表示一个定时消息条目。
     */
    @Getter
    @AllArgsConstructor
    public static class TimerLogEntry {
        private final long deliveryTimestamp;
        private final long commitLogOffset;
        private final long physicalPosition; // 条目在文件中的物理位置
    }

    public TimerMessageStore(String basePath, int wheelSize, long segmentSize) throws IOException {
        this.storePath = basePath + File.separator + "timer";
        this.wheelSize = wheelSize;
        this.segmentSize = segmentSize;
        this.slotTable = new ConcurrentSkipListMap<>();

        File dir = new File(this.storePath);
        if (!dir.exists() && !dir.mkdirs()) {
            throw new IOException("Failed to create timer store directory: " + this.storePath);
        }
        recover();
    }

    private void recover() throws IOException {
        File dir = new File(storePath);
        File[] files = dir.listFiles((d, name) -> name.endsWith(".timerlog"));

        if (files == null || files.length == 0) {
            log.info("No existing timer log found. Created initial segment.");
            this.activeSegment = new TimerLogSegment(new File(storePath, String.format("%020d.timerlog", 0L)).getAbsolutePath(), 0L, segmentSize);
            this.segments.put(0L, this.activeSegment);
            return;
        }

        for (File file : files) {
            long baseOffset = Long.parseLong(file.getName().replace(".timerlog", ""));
            TimerLogSegment segment = new TimerLogSegment(file.getAbsolutePath(), baseOffset, segmentSize);
            segments.put(baseOffset, segment);
            writePosition.set(baseOffset + segment.getWritePosition().get());
        }

        this.activeSegment = segments.lastEntry().getValue();

        // Rebuild slot table from the last entry of each slot
        log.info("Rebuilding slot table from recovered segments...");
        for (TimerLogSegment segment : segments.values()) {
            // 【修复】使用新的 getEntries 方法来恢复 slotTable
            List<TimerLogEntry> entries = segment.getEntries(0);
            for (TimerLogEntry entry : entries) {
                int slot = (int) ((entry.getDeliveryTimestamp() / 1000) % wheelSize);
                // 物理位置是段的基地址 + 在段内的相对位置
                long absolutePosition = segment.getBaseOffset() + entry.getPhysicalPosition();
                slotTable.put(slot, (int) absolutePosition);
            }
        }
        log.info("TimerMessageStore recovered. Total segments: {}. Next write position: {}. Slot table size: {}", segments.size(), writePosition.get(), slotTable.size());
    }

    public void enqueue(Message message, long commitLogOffset) throws IOException {
        long delaySeconds = message.getDelayTimeInSeconds();
        if (delaySeconds <= 0) {
            return;
        }
        long deliveryTimestamp = System.currentTimeMillis() + delaySeconds * 1000;
        int slot = (int) ((deliveryTimestamp / 1000) % this.wheelSize);

        lock.lock();
        try {
            // 【修复】调用 isFull() 时不再需要参数
            if (activeSegment.isFull()) {
                roll();
            }
            // 【修复】调用 appendEntry 时只传递需要的两个参数
            if (activeSegment.appendEntry(deliveryTimestamp, commitLogOffset)) {
                long absolutePosition = activeSegment.getBaseOffset() + activeSegment.getWritePosition().get() - TimerLogSegment.TIMER_LOG_ENTRY_SIZE;
                slotTable.put(slot, (int) absolutePosition);
                writePosition.addAndGet(TimerLogSegment.TIMER_LOG_ENTRY_SIZE);
            } else {
                // isFull() 可能有竞争条件，如果追加失败，再次尝试滚动
                roll();
                // 【修复】再次调用 isFull() 时不再需要参数
                if (activeSegment.isFull()) {
                    throw new IOException("Timer log is full, cannot roll to a new segment.");
                }
                // 【修复】再次调用 appendEntry
                if (activeSegment.appendEntry(deliveryTimestamp, commitLogOffset)) {
                    long absolutePosition = activeSegment.getBaseOffset() + activeSegment.getWritePosition().get() - TimerLogSegment.TIMER_LOG_ENTRY_SIZE;
                    slotTable.put(slot, (int) absolutePosition);
                    writePosition.addAndGet(TimerLogSegment.TIMER_LOG_ENTRY_SIZE);
                } else {
                    throw new IOException("Failed to append to new timer log segment.");
                }
            }
        } finally {
            lock.unlock();
        }
    }

    public List<TimerLogEntry> fetchAndClearSlotEntries(int slot) {
        Integer physicalPosition = this.slotTable.remove(slot);
        if (physicalPosition == null) {
            return Collections.emptyList();
        }

        Map.Entry<Long, TimerLogSegment> segmentEntry = this.segments.floorEntry((long) physicalPosition);
        if (segmentEntry == null) {
            log.warn("Cannot find segment for physical position: {}", physicalPosition);
            return Collections.emptyList();
        }

        TimerLogSegment segment = segmentEntry.getValue();
        // 【修复】计算在段内的相对位置
        int positionInSegment = (int) (physicalPosition - segment.getBaseOffset());

        // 【修复】使用新的 getEntries 方法
        List<TimerLogEntry> entries = segment.getEntries(positionInSegment);

        if (!entries.isEmpty()) {
            // 【修复】使用新的 clearEntries 方法
            segment.clearEntries(positionInSegment, entries.size());
        }

        return entries;
    }

    private void roll() throws IOException {
        long newBaseOffset = activeSegment.getBaseOffset() + segmentSize;
        File newFile = new File(storePath, String.format("%020d.timerlog", newBaseOffset));
        TimerLogSegment newSegment = new TimerLogSegment(newFile.getAbsolutePath(), newBaseOffset, segmentSize);
        segments.put(newBaseOffset, newSegment);
        activeSegment = newSegment;
        log.info("Rolled to new timer log segment: {}", newFile.getName());
    }

    @Override
    public void close() throws IOException {
        log.info("Closing TimerMessageStore...");
        lock.lock();
        try {
            for (TimerLogSegment segment : segments.values()) {
                segment.close();
            }
            segments.clear();
        } finally {
            lock.unlock();
            log.info("TimerMessageStore closed.");
        }
    }
}